Skip to content

[SPARK-59400][ML] Reuse TreeEnsembleModel prediction helper - #58691

Closed
zhengruifeng wants to merge 7 commits into
apache:masterfrom
zhengruifeng:ml-reuse-tree-ensemble-predict-dev-6
Closed

[SPARK-59400][ML] Reuse TreeEnsembleModel prediction helper#58691
zhengruifeng wants to merge 7 commits into
apache:masterfrom
zhengruifeng:ml-reuse-tree-ensemble-predict-dev-6

Conversation

@zhengruifeng

@zhengruifeng zhengruifeng commented Sep 10, 2026

Copy link
Copy Markdown
Contributor

What changes were proposed in this pull request?

This pull request centralizes tree ensemble prediction logic in companion objects:

  • Add weighted and unweighted TreeEnsembleModel.predictRaw overloads that accept tree models.
  • Rename the existing root-node helper to predictRaw for consistent naming.
  • Reuse these helpers in GBT classification, GBT regression, and random forest regression.
  • Move random forest classification's tree-based predictRaw implementation to its companion
    object and delegate to it from the model.

Why are the changes needed?

The model classes duplicated tree traversal and aggregation loops. Keeping the implementations in
companion-object helpers makes the direct model prediction paths consistent with the optimized
column prediction paths. The unweighted overload also avoids allocating an intermediate array when
random forest regression sums tree predictions.

Does this PR introduce any user-facing change?

No.

How was this patch tested?

No new tests were added because this is a behavior-preserving refactor of prediction paths already
covered by the existing GBT and random forest suites.

The patch passed git diff --check, changed-file line-length checks, and changed-file non-ASCII
checks. Unit tests were not run locally.

Was this patch authored or co-authored using generative AI tooling?

Generated-by: OpenAI Codex (GPT-5)

zhengruifeng added a commit that referenced this pull request Sep 11, 2026
### What changes were proposed in this pull request?

This pull request centralizes tree ensemble prediction logic in companion objects:

- Add weighted and unweighted `TreeEnsembleModel.predictRaw` overloads that accept tree models.
- Rename the existing root-node helper to `predictRaw` for consistent naming.
- Reuse these helpers in GBT classification, GBT regression, and random forest regression.
- Move random forest classification's tree-based `predictRaw` implementation to its companion
  object and delegate to it from the model.

### Why are the changes needed?

The model classes duplicated tree traversal and aggregation loops. Keeping the implementations in
companion-object helpers makes the direct model prediction paths consistent with the optimized
column prediction paths. The unweighted overload also avoids allocating an intermediate array when
random forest regression sums tree predictions.

### Does this PR introduce _any_ user-facing change?

No.

### How was this patch tested?

No new tests were added because this is a behavior-preserving refactor of prediction paths already
covered by the existing GBT and random forest suites.

The patch passed `git diff --check`, changed-file line-length checks, and changed-file non-ASCII
checks. Unit tests were not run locally.

### Was this patch authored or co-authored using generative AI tooling?

Generated-by: OpenAI Codex (GPT-5)

Closes #58691 from zhengruifeng/ml-reuse-tree-ensemble-predict-dev-6.

Authored-by: Ruifeng Zheng <ruifengz@apache.org>
Signed-off-by: Ruifeng Zheng <ruifengz@foxmail.com>
(cherry picked from commit 619f4f9)
Signed-off-by: Ruifeng Zheng <ruifengz@foxmail.com>
@zhengruifeng

Copy link
Copy Markdown
Contributor Author

Merge Summary:

Posted by merge_spark_pr.py

@zhengruifeng
zhengruifeng deleted the ml-reuse-tree-ensemble-predict-dev-6 branch September 11, 2026 00:13
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants